#!/usr/bin/env python3
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class GridSampleModule(nn.Module):
    def __init__(self):
        super(GridSampleModule, self).__init__()
    
    def forward(self, input_tensor, grid):
        # 使用torch的grid_sample函数
        return F.grid_sample(input_tensor, grid, mode='bilinear', 
                           padding_mode='zeros', align_corners=False)

def create_gridsample_onnx():
    print("=== 生成GridSample ONNX模型 ===")
    
    # 创建模型
    model = GridSampleModule()
    model.eval()
    
    # 定义输入尺寸
    batch_size = 1
    channels = 3
    height = 8
    width = 8
    
    # 创建输入张量
    input_tensor = torch.randn(batch_size, channels, height, width, dtype=torch.float32)
    
    # 创建grid张量 (归一化坐标)
    grid_h, grid_w = 6, 6
    grid = torch.randn(batch_size, grid_h, grid_w, 2, dtype=torch.float32)
    # 归一化到[-1, 1]范围
    grid = grid * 0.5
    
    print(f"输入张量形状: {input_tensor.shape}")
    print(f"Grid张量形状: {grid.shape}")
    
    # 测试前向传播
    with torch.no_grad():
        output = model(input_tensor, grid)
        print(f"输出张量形状: {output.shape}")
    
    # 导出ONNX
    onnx_path = "gridsample_test_model.onnx"
    
    # 创建示例输入用于ONNX导出
    dummy_input = torch.randn(batch_size, channels, height, width, dtype=torch.float32)
    dummy_grid = torch.randn(batch_size, grid_h, grid_w, 2, dtype=torch.float32) * 0.5
    
    print(f"导出ONNX模型到: {onnx_path}")
    
    torch.onnx.export(
        model,
        (dummy_input, dummy_grid),
        onnx_path,
        export_params=True,
        opset_version=16,  # 更新到16以支持grid_sampler
        do_constant_folding=True,
        input_names=['input', 'grid'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'grid': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        }
    )
    
    print("✅ ONNX模型导出成功!")
    
    # 验证ONNX模型
    import onnx
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    print("✅ ONNX模型验证通过!")
    
    # 打印模型信息
    print(f"ONNX IR版本: {onnx_model.ir_version}")
    print(f"Opset版本: {onnx_model.opset_import[0].version}")
    print(f"生产者: {onnx_model.producer_name}")
    
    print("输入:")
    for input_info in onnx_model.graph.input:
        shape = [dim.dim_value for dim in input_info.type.tensor_type.shape.dim]
        print(f"  {input_info.name}: {shape}")
    
    print("输出:")
    for output_info in onnx_model.graph.output:
        shape = [dim.dim_value for dim in output_info.type.tensor_type.shape.dim]
        print(f"  {output_info.name}: {shape}")
    
    return onnx_path

if __name__ == "__main__":
    create_gridsample_onnx() 